{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3b2e6ab7",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PreTrainedTokenizerFast(name_or_path='distilroberta-base', vocab_size=50265, model_max_len=512, is_fast=True, padding_side='right', truncation_side='right', special_tokens={'bos_token': '<s>', 'eos_token': '</s>', 'unk_token': '<unk>', 'sep_token': '</s>', 'pad_token': '<pad>', 'cls_token': '<s>', 'mask_token': AddedToken(\"<mask>\", rstrip=False, lstrip=True, single_word=False, normalized=False)})\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input_ids': [[0, 37265, 92, 3556, 2485, 31, 5, 20536, 2833, 2], [0, 10800, 5069, 117, 22094, 2156, 129, 6348, 3995, 821, 8299, 2]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "#加载编码器\n",
    "tokenizer = AutoTokenizer.from_pretrained('distilroberta-base', use_fast=True)\n",
    "\n",
    "print(tokenizer)\n",
    "\n",
    "#编码试算\n",
    "tokenizer.batch_encode_plus([\n",
    "    'hide new secretions from the parental units',\n",
    "    'contains no wit , only labored gags'\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "69c480e7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-23e50d2eeef7f0ac.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-70d4bf21e64d6e79.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-ee12ec06b02ec0f0.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-3907c59016f2be66.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-d519ea4f908e2cc9.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-4f9a17255014b2fe.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-cc742c0ee02005a1.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-b72c5d7ee8d2661a.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-3b4eaf100efe68a3.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-558c7e7de9aa763e.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-a4614e2dbcb4360e.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-20f4a0585ba9a4d2.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-92673eaaece2341d_00000_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-92673eaaece2341d_00001_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-92673eaaece2341d_00002_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-92673eaaece2341d_00003_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-39fb03207a78d634_00000_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-39fb03207a78d634_00001_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-39fb03207a78d634_00002_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-39fb03207a78d634_00003_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-e230ea8c3afd3870_00000_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-e230ea8c3afd3870_00001_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-e230ea8c3afd3870_00002_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-e230ea8c3afd3870_00003_of_00004.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-69f4464365084cd1.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-0a5268c6987f06db.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-9cdbc155b4fe7f84.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/train/cache-4c0e93b31b1854e5.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-7d3e669c2812fdc5.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-513b125c419d9621.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-4202180e45381dff.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/validation/cache-926b7c1917c3b05d.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-2de28716ffa0d538.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-f885a0014aa8f325.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-0a6008bd29f57cab.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at datas/glue/sst2/test/cache-04b3fc49d2f2242f.arrow\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(DatasetDict({\n",
       "     train: Dataset({\n",
       "         features: ['input_ids', 'attention_mask', 'labels'],\n",
       "         num_rows: 44279\n",
       "     })\n",
       "     validation: Dataset({\n",
       "         features: ['input_ids', 'attention_mask', 'labels'],\n",
       "         num_rows: 861\n",
       "     })\n",
       "     test: Dataset({\n",
       "         features: ['input_ids', 'attention_mask', 'labels'],\n",
       "         num_rows: 1776\n",
       "     })\n",
       " }),\n",
       " {'input_ids': [0, 37265, 92, 3556, 50264, 31, 5, 20536, 2],\n",
       "  'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "  'labels': [-100, -100, -100, -100, 2485, -100, -100, -100, -100]})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset, load_from_disk\n",
    "\n",
    "#加载数据\n",
    "#dataset = load_dataset(path='glue', name='sst2')\n",
    "dataset = load_from_disk('datas/glue/sst2')\n",
    "\n",
    "\n",
    "#分词,同时删除多余的字段\n",
    "def f(data):\n",
    "    return tokenizer.batch_encode_plus(data['sentence'])\n",
    "\n",
    "\n",
    "dataset = dataset.map(f,\n",
    "                      batched=True,\n",
    "                      batch_size=1000,\n",
    "                      num_proc=4,\n",
    "                      remove_columns=['sentence', 'idx', 'label'])\n",
    "\n",
    "\n",
    "#过滤掉太短的句子\n",
    "def f(data):\n",
    "    return [len(i) >= 9 for i in data['input_ids']]\n",
    "\n",
    "\n",
    "dataset = dataset.filter(f, batched=True, batch_size=1000, num_proc=4)\n",
    "\n",
    "\n",
    "#截断句子,同时整理成模型需要的格式\n",
    "def f(data):\n",
    "    b = len(data['input_ids'])\n",
    "    data['labels'] = data['attention_mask'].copy()\n",
    "    for i in range(b):\n",
    "        #裁剪长度到9\n",
    "        data['input_ids'][i] = data['input_ids'][i][:9]\n",
    "        data['attention_mask'][i] = [1] * 9\n",
    "        data['labels'][i] = [-100] * 9\n",
    "\n",
    "        #input_ids最后一位是2\n",
    "        data['input_ids'][i][-1] = 2\n",
    "\n",
    "        #每一句话第4个词为mask\n",
    "        #tokenizer.get_vocab()['<mask>'] -> 50264\n",
    "        data['labels'][i][4] = data['input_ids'][i][4]\n",
    "        data['input_ids'][i][4] = 50264\n",
    "\n",
    "    return data\n",
    "\n",
    "\n",
    "dataset = dataset.map(f, batched=True, batch_size=1000, num_proc=4)\n",
    "\n",
    "dataset, dataset['train'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ca303846",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5534,\n",
       " {'input_ids': tensor([[    0, 12905,     5,   247, 50264,   128,   197,   393,     2],\n",
       "          [    0,  1610,    10,   144, 50264,    12, 18948,   621,     2],\n",
       "          [    0,  7109, 35138,  4504, 50264,    15,   372, 32327,     2],\n",
       "          [    0,   212, 30990,  5475, 50264,    65,   631,    16,     2],\n",
       "          [    0, 23428,   149,  2156, 50264,    15,    10,  1402,     2],\n",
       "          [    0,  6025, 18013,  6629, 50264,    15,     5, 18754,     2],\n",
       "          [    0,  2629, 45518,  1462, 50264, 17957,    31,  1684,     2],\n",
       "          [    0,    90,  4894,   615, 50264,  2564,    11,   143,     2]]),\n",
       "  'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "          [1, 1, 1, 1, 1, 1, 1, 1, 1]]),\n",
       "  'labels': tensor([[ -100,  -100,  -100,  -100, 10361,  -100,  -100,  -100,  -100],\n",
       "          [ -100,  -100,  -100,  -100,   543,  -100,  -100,  -100,  -100],\n",
       "          [ -100,  -100,  -100,  -100,  1109,  -100,  -100,  -100,  -100],\n",
       "          [ -100,  -100,  -100,  -100,    59,  -100,  -100,  -100,  -100],\n",
       "          [ -100,  -100,  -100,  -100,  2254,  -100,  -100,  -100,  -100],\n",
       "          [ -100,  -100,  -100,  -100,    62,  -100,  -100,  -100,  -100],\n",
       "          [ -100,  -100,  -100,  -100,  1141,  -100,  -100,  -100,  -100],\n",
       "          [ -100,  -100,  -100,  -100,     7,  -100,  -100,  -100,  -100]])})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers.data.data_collator import default_data_collator\n",
    "\n",
    "#能够实现随机mask的collate_fn\n",
    "#如果要使用这个工具类,在数据预处理时就不需要设置数据中的mask,然后让labels=input_ids.copy即可\n",
    "#from transformers import DataCollatorForLanguageModeling\n",
    "#data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer,mlm_probability=0.1)\n",
    "\n",
    "#数据加载器\n",
    "loader = torch.utils.data.DataLoader(\n",
    "    dataset=dataset['train'],\n",
    "    batch_size=8,\n",
    "    collate_fn=default_data_collator,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    ")\n",
    "\n",
    "for i, data in enumerate(loader):\n",
    "    break\n",
    "\n",
    "len(loader), data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d9488731",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at distilroberta-base were not used when initializing RobertaModel: ['lm_head.bias', 'lm_head.layer_norm.weight', 'lm_head.dense.bias', 'lm_head.dense.weight', 'lm_head.decoder.weight', 'lm_head.layer_norm.bias']\n",
      "- This IS expected if you are initializing RobertaModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing RobertaModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12136.4313\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(tensor(19.1579, grad_fn=<NllLossBackward0>), torch.Size([8, 9, 50265]))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoModelForCausalLM, RobertaModel\n",
    "\n",
    "#加载模型\n",
    "#model = AutoModelForCausalLM.from_pretrained('distilroberta-base')\n",
    "\n",
    "\n",
    "#定义下游任务模型\n",
    "class Model(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.pretrained = RobertaModel.from_pretrained('distilroberta-base')\n",
    "\n",
    "        decoder = torch.nn.Linear(768, tokenizer.vocab_size)\n",
    "        decoder.bias = torch.nn.Parameter(torch.zeros(tokenizer.vocab_size))\n",
    "\n",
    "        self.fc = torch.nn.Sequential(\n",
    "            torch.nn.Linear(768, 768),\n",
    "            torch.nn.GELU(),\n",
    "            torch.nn.LayerNorm(768, eps=1e-5),\n",
    "            decoder,\n",
    "        )\n",
    "\n",
    "        #加载预训练模型的参数\n",
    "        parameters = AutoModelForCausalLM.from_pretrained('distilroberta-base')\n",
    "        self.fc[0].load_state_dict(parameters.lm_head.dense.state_dict())\n",
    "        self.fc[2].load_state_dict(parameters.lm_head.layer_norm.state_dict())\n",
    "        self.fc[3].load_state_dict(parameters.lm_head.decoder.state_dict())\n",
    "\n",
    "        self.criterion = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, input_ids, attention_mask, labels=None):\n",
    "        logits = self.pretrained(input_ids=input_ids,\n",
    "                                 attention_mask=attention_mask)\n",
    "        logits = logits.last_hidden_state\n",
    "\n",
    "        logits = self.fc(logits)\n",
    "\n",
    "        loss = None\n",
    "        if labels is not None:\n",
    "            shifted_logits = logits[:, :-1].reshape(-1, tokenizer.vocab_size)\n",
    "            shifted_labels = labels[:, 1:].reshape(-1)\n",
    "\n",
    "            loss = self.criterion(shifted_logits, shifted_labels)\n",
    "\n",
    "        return {'loss': loss, 'logits': logits}\n",
    "\n",
    "\n",
    "model = Model()\n",
    "\n",
    "#统计参数量\n",
    "print(sum(i.numel() for i in model.parameters()) / 10000)\n",
    "\n",
    "out = model(**data)\n",
    "\n",
    "out['loss'], out['logits'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c924a315",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "tensor([   47, 14838,  5392,    28,    80,  4839,  3668,    29])\n",
      "tensor([   47, 14633,   749,    28,    80,  4839,  3668,  2156])\n",
      "10\n",
      "tensor([ 101,  668,   16,   14,  352,  650, 3961,   16])\n",
      "tensor([ 101,  773, 7897,   59, 2156, 7397, 3961,   16])\n",
      "20\n",
      "tensor([40485,    13,    29, 19303,    33,    16,   295,     9])\n",
      "tensor([40485,    13,  4839, 16393,    33,  3391,   256,     9])\n",
      "30\n",
      "tensor([   53, 33469,  3315,  3723,     7, 24473, 40776,    41])\n",
      "tensor([11248, 15923,  3315,  3723,     7, 24473, 40776,    41])\n",
      "40\n",
      "tensor([ 2435,     5,  2046,  2084, 25210,     9, 42661,     7])\n",
      "tensor([ 2343,    42,  4265,  8003, 33709,  7021,  9021,     6])\n",
      "50\n",
      "tensor([  297, 22258,   998,    64,    10,  1499,    65,  2156])\n",
      "tensor([  457, 22258,  6545,    64,    10, 10416,    65, 33647])\n",
      "0.32598039215686275\n",
      "<s>a strong first<mask>, slightly less</s>\n",
      " quarter  half\n",
      "<s>( villene<mask> ) seems to</s>\n",
      "uve uve\n",
      "<s>going to the<mask> may be just</s>\n",
      " website  gym\n",
      "<s>anyone who<mask> count to five</s>\n",
      " can  can\n",
      "<s>it's<mask> terrific american</s>\n",
      " a  a\n",
      "<s>this is car<mask>'s debut</s>\n",
      "ion pool\n",
      "<s>the movie is<mask> of the best</s>\n",
      " one  one\n",
      "<s>how on earth<mask> or anywhere else</s>\n",
      ", ?,\n"
     ]
    }
   ],
   "source": [
    "#测试\n",
    "def test():\n",
    "    model.eval()\n",
    "\n",
    "    loader_test = torch.utils.data.DataLoader(\n",
    "        dataset=dataset['test'],\n",
    "        batch_size=8,\n",
    "        collate_fn=default_data_collator,\n",
    "        shuffle=True,\n",
    "        drop_last=True,\n",
    "    )\n",
    "\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    for i, data in enumerate(loader_test):\n",
    "\n",
    "        #保存下数据中的label,后面计算正确率要用\n",
    "        label = data['labels'][:, 4].clone()\n",
    "\n",
    "        #从数据中抹除掉label,防止模型作弊\n",
    "        data['labels'] = None\n",
    "\n",
    "        #计算\n",
    "        with torch.no_grad():\n",
    "            out = model(**data)\n",
    "\n",
    "        #[8, 10, 50265] -> [8, 10]\n",
    "        out = out['logits'].argmax(dim=2)[:, 4]\n",
    "\n",
    "        correct += (label == out).sum().item()\n",
    "        total += 8\n",
    "\n",
    "        if i % 10 == 0:\n",
    "            print(i)\n",
    "            print(label)\n",
    "            print(out)\n",
    "\n",
    "        if i == 50:\n",
    "            break\n",
    "\n",
    "    print(correct / total)\n",
    "\n",
    "    for i in range(8):\n",
    "        print(tokenizer.decode(data['input_ids'][i]))\n",
    "        print(tokenizer.decode(label[i]), tokenizer.decode(out[i]))\n",
    "\n",
    "\n",
    "test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1375fcad",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/cpu/lib/python3.6/site-packages/transformers/optimization.py:309: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  FutureWarning,\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 18.949838638305664 0.0 1.9996385977593064e-05\n",
      "50 4.755198001861572 0.625 1.9815684857246115e-05\n",
      "100 5.0272216796875 0.25 1.963498373689917e-05\n",
      "150 4.625316143035889 0.125 1.9454282616552225e-05\n",
      "200 3.663780927658081 0.5 1.927358149620528e-05\n",
      "250 2.5342917442321777 0.375 1.909288037585833e-05\n",
      "300 4.986537933349609 0.375 1.8912179255511386e-05\n",
      "350 3.403028964996338 0.625 1.873147813516444e-05\n",
      "400 4.041268348693848 0.125 1.8550777014817495e-05\n",
      "450 3.2715964317321777 0.5 1.8370075894470547e-05\n",
      "500 2.6591811180114746 0.5 1.81893747741236e-05\n",
      "550 4.937175750732422 0.25 1.8008673653776656e-05\n",
      "600 4.845945835113525 0.25 1.7827972533429708e-05\n",
      "650 1.8658218383789062 0.625 1.7647271413082763e-05\n",
      "700 3.9473319053649902 0.25 1.7466570292735818e-05\n",
      "750 2.065851926803589 0.625 1.728586917238887e-05\n",
      "800 2.957096576690674 0.5 1.7105168052041924e-05\n",
      "850 4.987250804901123 0.25 1.692446693169498e-05\n",
      "900 3.5697021484375 0.5 1.674376581134803e-05\n",
      "950 2.898092746734619 0.5 1.6563064691001085e-05\n",
      "1000 4.39031457901001 0.375 1.638236357065414e-05\n",
      "1050 4.365252494812012 0.5 1.6201662450307194e-05\n",
      "1100 3.3205385208129883 0.625 1.6020961329960246e-05\n",
      "1150 3.8429198265075684 0.5 1.58402602096133e-05\n",
      "1200 1.5565372705459595 0.5 1.5659559089266356e-05\n",
      "1250 3.046530246734619 0.375 1.547885796891941e-05\n",
      "1300 1.1127394437789917 0.625 1.5298156848572462e-05\n",
      "1350 3.349091053009033 0.5 1.5117455728225517e-05\n",
      "1400 3.6459503173828125 0.25 1.493675460787857e-05\n",
      "1450 1.958971619606018 0.625 1.4756053487531623e-05\n",
      "1500 1.8479055166244507 0.5 1.457535236718468e-05\n",
      "1550 2.961282730102539 0.375 1.4394651246837732e-05\n",
      "1600 1.8593209981918335 0.625 1.4213950126490786e-05\n",
      "1650 3.1550772190093994 0.5 1.4033249006143839e-05\n",
      "1700 0.8767207860946655 0.5 1.3852547885796892e-05\n",
      "1750 1.7590525150299072 0.625 1.3671846765449948e-05\n",
      "1800 1.7380354404449463 0.625 1.3491145645103002e-05\n",
      "1850 2.374969959259033 0.5 1.3310444524756055e-05\n",
      "1900 2.9406280517578125 0.375 1.3129743404409108e-05\n",
      "1950 1.2378032207489014 1.0 1.2949042284062163e-05\n",
      "2000 2.7496113777160645 0.5 1.2768341163715217e-05\n",
      "2050 2.0767924785614014 0.75 1.258764004336827e-05\n",
      "2100 2.7222516536712646 0.5 1.2406938923021324e-05\n",
      "2150 3.213247537612915 0.5 1.2226237802674377e-05\n",
      "2200 3.55709171295166 0.25 1.2045536682327432e-05\n",
      "2250 1.850517988204956 0.5 1.1864835561980486e-05\n",
      "2300 2.4408578872680664 0.5 1.168413444163354e-05\n",
      "2350 1.7818212509155273 0.5 1.1503433321286593e-05\n",
      "2400 1.1785320043563843 0.625 1.1322732200939646e-05\n",
      "2450 2.774081230163574 0.625 1.11420310805927e-05\n",
      "2500 2.8569412231445312 0.5 1.0961329960245755e-05\n",
      "2550 0.8837183713912964 0.875 1.0780628839898809e-05\n",
      "2600 1.9974981546401978 0.5 1.0599927719551862e-05\n",
      "2650 3.6552603244781494 0.375 1.0419226599204916e-05\n",
      "2700 3.2599399089813232 0.375 1.023852547885797e-05\n",
      "2750 2.019043445587158 0.625 1.0057824358511023e-05\n",
      "2800 1.3741586208343506 0.625 9.877123238164078e-06\n",
      "2850 1.3488343954086304 0.75 9.69642211781713e-06\n",
      "2900 2.849611520767212 0.375 9.515720997470185e-06\n",
      "2950 2.4670894145965576 0.5 9.335019877123239e-06\n",
      "3000 2.3208529949188232 0.625 9.154318756776292e-06\n",
      "3050 1.7063318490982056 0.625 8.973617636429347e-06\n",
      "3100 3.4048428535461426 0.25 8.7929165160824e-06\n",
      "3150 1.6324763298034668 0.875 8.612215395735455e-06\n",
      "3200 4.0897088050842285 0.625 8.431514275388508e-06\n",
      "3250 5.042566299438477 0.375 8.25081315504156e-06\n",
      "3300 4.727066516876221 0.25 8.070112034694616e-06\n",
      "3350 1.72847318649292 0.75 7.889410914347669e-06\n",
      "3400 2.650611162185669 0.5 7.708709794000724e-06\n",
      "3450 1.8998345136642456 0.375 7.5280086736537775e-06\n",
      "3500 2.720449686050415 0.625 7.3473075533068315e-06\n",
      "3550 1.8739738464355469 0.625 7.166606432959885e-06\n",
      "3600 2.216492176055908 0.625 6.9859053126129386e-06\n",
      "3650 1.9913058280944824 0.5 6.8052041922659925e-06\n",
      "3700 2.00622296333313 0.5 6.6245030719190465e-06\n",
      "3750 4.7002434730529785 0.25 6.4438019515721005e-06\n",
      "3800 4.12710428237915 0.25 6.263100831225154e-06\n",
      "3850 0.9240655303001404 0.625 6.082399710878208e-06\n",
      "3900 2.604069709777832 0.375 5.9016985905312615e-06\n",
      "3950 1.720001220703125 0.375 5.7209974701843155e-06\n",
      "4000 1.4786216020584106 0.75 5.5402963498373695e-06\n",
      "4050 1.1555869579315186 0.75 5.359595229490423e-06\n",
      "4100 0.5870824456214905 0.875 5.1788941091434774e-06\n",
      "4150 1.1685823202133179 0.625 4.9981929887965306e-06\n",
      "4200 1.917506217956543 0.5 4.8174918684495845e-06\n",
      "4250 2.6694304943084717 0.375 4.6367907481026385e-06\n",
      "4300 4.763072967529297 0.25 4.4560896277556925e-06\n",
      "4350 2.5590553283691406 0.375 4.2753885074087464e-06\n",
      "4400 1.7727317810058594 0.5 4.0946873870618e-06\n",
      "4450 1.6883273124694824 0.75 3.9139862667148535e-06\n",
      "4500 2.8899121284484863 0.625 3.7332851463679075e-06\n",
      "4550 3.0935513973236084 0.375 3.5525840260209615e-06\n",
      "4600 1.5275040864944458 0.5 3.3718829056740154e-06\n",
      "4650 2.4337785243988037 0.5 3.1911817853270694e-06\n",
      "4700 1.1596354246139526 0.875 3.0104806649801234e-06\n",
      "4750 3.4586596488952637 0.375 2.8297795446331774e-06\n",
      "4800 1.3651186227798462 0.375 2.6490784242862305e-06\n",
      "4850 3.6087186336517334 0.25 2.468377303939285e-06\n",
      "4900 3.427321672439575 0.375 2.2876761835923384e-06\n",
      "4950 0.834938645362854 0.75 2.1069750632453924e-06\n",
      "5000 3.0450022220611572 0.5 1.926273942898446e-06\n",
      "5050 2.685086250305176 0.625 1.7455728225515001e-06\n",
      "5100 1.1113600730895996 0.5 1.5648717022045537e-06\n",
      "5150 1.7858939170837402 0.75 1.3841705818576077e-06\n",
      "5200 1.8797955513000488 0.625 1.2034694615106614e-06\n",
      "5250 2.4701387882232666 0.5 1.0227683411637154e-06\n",
      "5300 2.002549886703491 0.625 8.420672208167693e-07\n",
      "5350 1.2900426387786865 0.75 6.61366100469823e-07\n",
      "5400 1.4372508525848389 0.75 4.806649801228768e-07\n",
      "5450 1.6830159425735474 0.75 2.9996385977593064e-07\n",
      "5500 2.349522590637207 0.5 1.1926273942898448e-07\n"
     ]
    }
   ],
   "source": [
    "from transformers import AdamW\n",
    "from transformers.optimization import get_scheduler\n",
    "\n",
    "\n",
    "#训练\n",
    "def train():\n",
    "    optimizer = AdamW(model.parameters(), lr=2e-5)\n",
    "    scheduler = get_scheduler(name='linear',\n",
    "                              num_warmup_steps=0,\n",
    "                              num_training_steps=len(loader),\n",
    "                              optimizer=optimizer)\n",
    "\n",
    "    model.train()\n",
    "    for i, data in enumerate(loader):\n",
    "        out = model(**data)\n",
    "        loss = out['loss']\n",
    "\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        model.zero_grad()\n",
    "\n",
    "        if i % 50 == 0:\n",
    "            label = data['labels'][:, 4]\n",
    "            out = out['logits'].argmax(dim=2)[:, 4]\n",
    "\n",
    "            correct = (label == out).sum().item()\n",
    "            accuracy = correct / 8\n",
    "\n",
    "            lr = optimizer.state_dict()['param_groups'][0]['lr']\n",
    "\n",
    "            print(i, loss.item(), accuracy, lr)\n",
    "\n",
    "    torch.save(model, 'models/2.预测中间词.model')\n",
    "\n",
    "\n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ed397281",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "tensor([ 8444,   144,   543,     4,   253, 27229,  1071,  2230])\n",
      "tensor([ 5313, 33019,   543,     4,   283, 38821,   219, 41906])\n",
      "10\n",
      "tensor([ 3865, 31368,    84,    32,  1810,     9,     7,    24])\n",
      "tensor([3865, 6269,   84,   32, 1810,   19,    7,   24])\n",
      "20\n",
      "tensor([ 2156,  2156,  2156,   128,    29, 16887,  7458,    16])\n",
      "tensor([2156, 2156, 2156,  128,   29, 9669, 7458,   16])\n",
      "30\n",
      "tensor([43578,   475,   814,     5,   939,  2156,   266,  1766])\n",
      "tensor([2653,  475,  814,    5,   10, 2156,  192, 1766])\n",
      "40\n",
      "tensor([   12, 36302,    95, 19987,    16,  6476,   275,    32])\n",
      "tensor([   12,  6269,   215, 22019,  2156,  2156, 14535,    70])\n",
      "50\n",
      "tensor([ 1495,     8,  2156,  2156,  3025, 13760,  2789,   408])\n",
      "tensor([   95,  2156,  2156,     5,  3025,    34, 21468, 13670])\n",
      "0.49264705882352944\n",
      "<s>the film is<mask> a sort of</s>\n",
      " itself  just\n",
      "<s>there are slow<mask> repetitive parts,</s>\n",
      " and ,\n",
      "<s>theology aside<mask> why put someone</s>\n",
      ", ,\n",
      "<s>all ends well<mask> sort of,</s>\n",
      ",  the\n",
      "<s>shot perhaps `<mask>ically'with</s>\n",
      " artist  artist\n",
      "<s>oedek<mask> wrote patch ad</s>\n",
      "erk  has\n",
      "<s>showtime is<mask> to slowtime</s>\n",
      " closer  akin\n",
      "<s>will grab your<mask> by the imagination</s>\n",
      " children  imagination\n"
     ]
    }
   ],
   "source": [
    "model = torch.load('models/2.预测中间词.model')\n",
    "test()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
